/*
 * ManyTrees.cpp
 *
 *  Created on: Mar 22, 2012
 *      Author: pickrell
 */
#include "ManyTrees.h"
using namespace std;


ManyTrees::ManyTrees(){

}

ManyTrees::ManyTrees(string infile, PhyloPop_params *p ){

	// input file is location of alignment, location of tree

	d.clear();
	igzstream in(infile.c_str());
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	intStat = stat(infile.c_str(), &stFileInfo);
	if (intStat !=0){
         std::cerr<< "ERROR: cannot open file " << infile << "\n";
         exit(1);
	}
	string st;
	while(getline(in, st)){
		string buf;
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		Alignment aln(line[0]);
		TreeState tmp(&aln, p);
		tmp.set_to_newick(line[1]);
		d.push_back(tmp);
	}
	epsilon = 1e-10;
    phi = (1+sqrt(5))/2;
    resphi = 2-phi;
}


void ManyTrees::print(){
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++) it->print();
}

void ManyTrees::print_allcoaltimes(){
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++) it->print_coales();
}


vector<double> ManyTrees::get_all_coaltimes(){
	vector<double> toreturn;
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++){
		for (vector<double>::iterator it2 = it->coaltimes.begin(); it2 != it->coaltimes.end(); it2++){
			toreturn.push_back(*it2);
		}
	}
	sort(toreturn.begin(), toreturn.end());
	return toreturn;
}

void ManyTrees::set_bpoints(int N){
	bpoints.clear();
	vector<double> alltimes = get_all_coaltimes();
	for (int i = 1; i<= N; i++){
		double tmp = (double)i /(double)N;
		int which =  tmp * alltimes.size();
		which = which -1;
		//cout << i << " "<< tmp << " "<< alltimes.size() << " "<< which << " "<< alltimes[which] << "\n";
		double test = alltimes[which]+epsilon;
		if (bpoints.size() > 0 && fabs(test-bpoints[bpoints.size()-1])< 1e-10) continue;
		else if (fabs(test-epsilon) < 1e-10) continue;
		cout << i << " "<< tmp << " "<< alltimes.size() << " "<< which << " "<< alltimes[which] << "\n";
		if (i == N) bpoints.push_back(alltimes[which] + epsilon);
		else bpoints.push_back(alltimes[which] +epsilon);
	}
}

void ManyTrees::print_bpoints(){
	for (vector<double>::iterator it = bpoints.begin(); it != bpoints.end(); it++) cout <<*it << "\n";

}

void ManyTrees::init_popsizes(){
	popsizes.clear();
	for (int i = 0; i < bpoints.size(); i++){
		popsizes.push_back(1);
	}
}

void ManyTrees::init_structure(int which, gsl_rng * r){
	set_structured_bpoints_popsizes(which);
	cout << "t: "<< t << "\n";
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++){
		it->init_tip_labels(t, r);
	}
	setp1();
	//cout << "init p1: "<< p1 << "\n";
}

void ManyTrees::optim_all_struct(){
	optimize_all_tips();
	setp1();
	optim_all_struct_popsizes();

}

void ManyTrees::iterate_optim_all_struct(){
	double st_llk = llk_structured();
	bool go = true;
	while (go){
		optim_all_struct();
		double test_llk = llk_structured();
		cout << "start: "<< st_llk << " test: "<< test_llk << "\n";
		if (test_llk < st_llk + 0.001) go = false;
		else st_llk = test_llk;
	}
}

double ManyTrees::llk(){
	double toreturn = 0;
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++){
		//cout << "here "<< toreturn << "\n"; cout.flush();
		toreturn += it->llk(&bpoints, &popsizes);
		//cout << "ll1 "<< it->llk(&bpoints, &popsizes) << "\n";
		//cout << "done1\n"; cout.flush();
		//toreturn += it->llk(it->current_nhap, 0, 0, it->coaltimes, &bpoints, &popsizes);
		//cout << "ll2 "<< it->llk(it->current_nhap, 0, 0, it->coaltimes, &bpoints, &popsizes) << "\n";
		//cout << "done2\n"; cout.flush();
	}
	return toreturn;
}

void ManyTrees::print_popsizes(){
	for (vector<double>::iterator it = popsizes.begin(); it != popsizes.end(); it++){
		cout << *it << "\n";
	}
}

double ManyTrees::llk_structured(){
	double toreturn =0;
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++){
		toreturn+= it->structured_llk(t, p1, &bpoints1, &bpoints3, &popsizes1, &popsizes2, &popsizes3);
	}
	return toreturn;
}

double ManyTrees::llk_structured_no3(int which){
	return d.at(which).structured_llk_no3(p1, &bpoints1, &popsizes1, &popsizes2);
}
double ManyTrees::llk_structured_no3(){
	double toreturn =0;
	int i  =0;
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++){
		double toadd = it->structured_llk_no3(p1, &bpoints1, &popsizes1, &popsizes2);
		//cout << i << " "<< toadd << "\n";
		toreturn+= toadd;
		i++;
	}
	return toreturn;
}

void ManyTrees::setp1(){
	double tmp = 0;
	for (int i = 0; i < d.size(); i++) tmp += d.at(i).getp();
	tmp = tmp/ (double) d.size();
	p1 = tmp;
	cout << "new p1: "<< p1 << "\n";
}

void ManyTrees::set_structured_bpoints_popsizes(int which){
	if (which >=bpoints.size()){
		cerr<< "ERROR: trying to set structured breakpoint at "<< which << ", vector of size "<< bpoints.size()<< "\n";
		exit(1);
	}
	t = bpoints[which];
	bpoints1.clear();
	bpoints2.clear();
	bpoints3.clear();
	popsizes1.clear();
	popsizes2.clear();
	popsizes3.clear();
	for (int i = 0; i <= which; i++){
		bpoints1.push_back(bpoints[i]);
		bpoints2.push_back(bpoints[i]);
		popsizes1.push_back(popsizes[i]);
		popsizes2.push_back(popsizes[i]/2.0);
	}
	for (int i = which+1; i < bpoints.size(); i++){
		bpoints3.push_back(bpoints[i]);
		popsizes3.push_back(popsizes[i]);
	}
	bpoints3[bpoints3.size()-1] += 1e-5;

}

void ManyTrees::set_known(int which){
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++){
		it->set_known(which);
	}
}

void ManyTrees::optimize_all_tips(){
	for (vector<TreeState>::iterator it = d.begin(); it != d.end(); it++){
		it->iterate_optimize_all_tips(p1, &bpoints1, &popsizes1, &popsizes2);
	}
}

void ManyTrees::optim_struct_popsize(int whichbin, int whichpop){
	int nit = 0;
	double min = -10;
	double max = 10;
	double guess;
	if (whichpop ==1) guess = log(popsizes1[whichbin]);
	else if (whichpop ==2) guess = log(popsizes2[whichbin]);
	double tau = 0.001;
	golden_section_struct_popsize(whichbin, whichpop, min, guess, max, tau, &nit);
}


void ManyTrees::optim_all_struct_popsizes(){
	double old_llk = llk_structured_no3();
	for (int i = 0; i < popsizes1.size(); i++) {
		optim_struct_popsize(i, 1);
		optim_struct_popsize(i, 2);
	}
	double current_llk = llk_structured_no3();
	cout << "old :"<< old_llk << " current: "<< current_llk << "\n";
	while ( (current_llk - old_llk) > 1e-3 ){
		cout << "old :"<< old_llk << " current: "<< current_llk << "\n";
		old_llk = current_llk;
		for (int i = 0; i < popsizes1.size(); i++) {
			optim_struct_popsize(i, 1);
			optim_struct_popsize(i, 2);
		}
		current_llk = llk_structured_no3();
	}
}

int ManyTrees::golden_section_struct_popsize(int whichbin, int whichpop, double min, double guess, double max, double tau, int* nit){
        double x;

        //cout << guess << "\n"; cout.flush();
        if ( (max - guess) > (guess - min)) x = guess + resphi *( max - guess);
        else x = guess - resphi *(guess-min);
        if (fabs(max-min) < tau * (fabs(guess)+fabs(max))) {
                double new_logweight = (min+max)/2;
                double newsize = exp(new_logweight);
                if (whichpop ==1)popsizes1[whichbin] = newsize;
                if (whichpop ==2) popsizes2[whichbin] = newsize;
                //popsizes[which] = newsize;
                return 0;
        }
        *nit = *nit+1;
        double psize = exp(x);
        if (whichpop ==1) popsizes1[whichbin] = psize;
        if (whichpop ==2) popsizes2[whichbin] = psize;
        //popsizes[which] = psize;


        double f_x = -llk_structured_no3();


        psize = exp(guess);
        if (whichpop ==1) popsizes1[whichbin] = psize;
        if (whichpop ==2) popsizes2[whichbin] = psize;
        //popsizes[which] = psize;

        double f_guess = -llk_structured_no3();
        //cout << guess << " "<< -llik() << "\n";

        if (f_x < f_guess){
                if ( (max-guess) > (guess-min) ) return golden_section_struct_popsize(whichbin, whichpop, guess, x, max, tau, nit);
                else return golden_section_struct_popsize(whichbin, whichpop, min, x, guess, tau, nit);
        }

        else{
                if ( (max - guess) > (guess - min)  ) return golden_section_struct_popsize(whichbin, whichpop, min, guess, x, tau, nit);
                else return golden_section_struct_popsize(whichbin, whichpop, x, guess, max, tau, nit);
        }
}

void ManyTrees::print_state(string outfile){
	ofstream out(outfile.c_str());
	for (int i = 0; i < popsizes.size(); i++){
		double t1, t2, p;
		t2 = bpoints[i];
		if (i == 0 ) t1 = 0;
		else t1 = bpoints[i-1];
		p = popsizes[i];
		out << t1 << " "<< t2 << " "<< p <<"\n";
	}

}

void ManyTrees::print_all_trees(string outfile){
	for (int i = 0; i < d.size(); i++){
		stringstream ss;
		ss << outfile;
		ss<< ".tree";
		ss << i;
		d[i].tree->print(ss.str());
	}
}
void ManyTrees::print_struct_state(string outfile){
	ofstream out(outfile.c_str());
	for (int i = 0; i < popsizes1.size(); i++){
		double t1, t2, p;
		t2 = bpoints1[i];
		if (i == 0 ) t1 = 0;
		else t1 = bpoints1[i-1];
		p = popsizes1[i];
		out << "1 "<< t1 << " "<< t2 << " "<< p <<"\n";
	}
	for (int i = 0; i < popsizes2.size(); i++){
		double t1, t2, p;
		t2 = bpoints2[i];
		if (i == 0 ) t1 = 0;
		else t1 = bpoints2[i-1];
		p = popsizes2[i];
		out << "2 "<< t1 << " "<< t2 << " "<< p <<"\n";
	}
	for (int i = 0; i < popsizes3.size(); i++){
		double t1, t2, p;
		t2 = bpoints3[i];
		if (i == 0 ) t1 = t;
		else t1 = bpoints3[i-1];
		p = popsizes3[i];
		out << "3 "<< t1 << " "<< t2 << " "<< p <<"\n";
	}
}

void ManyTrees::optim_popsize(int which){
	int nit = 0;
	double min = -10;
	double max = 10;
	double guess = log(popsizes[which]);
	double tau = 0.001;
	golden_section_popsize(which, min, guess, max, tau, &nit);
}

void ManyTrees::optim_all_popsizes(){
	double old_llk = llk();
	for (int i = 0; i < popsizes.size(); i++) optim_popsize(i);
	double current_llk = llk();
	while ( (current_llk - old_llk) > 1e-3 ){
		old_llk = current_llk;
		for (int i = 0; i < popsizes.size(); i++) optim_popsize(i);
		current_llk = llk();
	}
}
int ManyTrees::golden_section_popsize(int which, double min, double guess, double max, double tau, int* nit){
        double x;

        //cout << guess << "\n"; cout.flush();
        if ( (max - guess) > (guess - min)) x = guess + resphi *( max - guess);
        else x = guess - resphi *(guess-min);
        if (fabs(max-min) < tau * (fabs(guess)+fabs(max))) {
                double new_logweight = (min+max)/2;
                double newsize = exp(new_logweight);
                popsizes[which] = newsize;
                return 0;
        }
        *nit = *nit+1;
        double psize = exp(x);
        popsizes[which] = psize;


        double f_x = -llk();


        psize = exp(guess);
        popsizes[which] = psize;

        double f_guess = -llk();
        //cout << guess << " "<< -llik() << "\n";

        if (f_x < f_guess){
                if ( (max-guess) > (guess-min) ) return golden_section_popsize(which, guess, x, max, tau, nit);
                else return golden_section_popsize(which, min, x, guess, tau, nit);
        }

        else{
                if ( (max - guess) > (guess - min)  ) return golden_section_popsize(which, min, guess, x, tau, nit);
                else return golden_section_popsize(which, x, guess, max, tau, nit);
        }
}

int ManyTrees::golden_section_popsize_bootstrap(vector<int> tokeep, int which, double min, double guess, double max, double tau, int* nit){
        double x;

        //cout << guess << "\n"; cout.flush();
        if ( (max - guess) > (guess - min)) x = guess + resphi *( max - guess);
        else x = guess - resphi *(guess-min);
        if (fabs(max-min) < tau * (fabs(guess)+fabs(max))) {
                double new_logweight = (min+max)/2;
                double newsize = exp(new_logweight);
                popsizes[which] = newsize;
                return 0;
        }
        *nit = *nit+1;
        double psize = exp(x);
        popsizes[which] = psize;


        double f_x = -llk_bootstrap(tokeep);


        psize = exp(guess);
        popsizes[which] = psize;

        double f_guess = -llk_bootstrap(tokeep);
        //cout << guess << " "<< -llik() << "\n";

        if (f_x < f_guess){
                if ( (max-guess) > (guess-min) ) return golden_section_popsize_bootstrap(tokeep, which, guess, x, max, tau, nit);
                else return golden_section_popsize_bootstrap(tokeep, which, min, x, guess, tau, nit);
        }

        else{
                if ( (max - guess) > (guess - min)  ) return golden_section_popsize_bootstrap(tokeep, which, min, guess, x, tau, nit);
                else return golden_section_popsize_bootstrap(tokeep, which, x, guess, max, tau, nit);
        }
}

void ManyTrees::set_bpoints(vector<double> toset){
	bpoints.clear();
	for (vector<double>::iterator it = toset.begin(); it != toset.end(); it++) bpoints.push_back(*it);
}

double ManyTrees::llk_bootstrap(vector<int> tokeep){
	double toreturn = 0;
	for (vector<int>::iterator it = tokeep.begin(); it != tokeep.end(); it++){
		toreturn += d[*it].llk(&bpoints, &popsizes);
	}
	return toreturn;

}
void ManyTrees::optim_all_popsizes_bootstrap(vector<int> tokeep){
	double old_llk = llk_bootstrap(tokeep);
	for (int i = 0; i < popsizes.size(); i++) optim_popsize_bootstrap(tokeep, i);
	double current_llk = llk_bootstrap(tokeep);
	while ( (current_llk - old_llk) > 1e-3 ){
		old_llk = current_llk;
		for (int i = 0; i < popsizes.size(); i++) optim_popsize_bootstrap(tokeep, i);
		current_llk = llk_bootstrap(tokeep);
	}
}

void ManyTrees::optim_popsize_bootstrap(vector<int> tokeep, int which){
	int nit = 0;
	double min = -10;
	double max = 10;
	double guess = log(popsizes[which]);
	double tau = 0.001;
	golden_section_popsize_bootstrap(tokeep, which, min, guess, max, tau, &nit);
}

vector<vector<double> > ManyTrees::bootstrap(int count, gsl_rng * r){
	vector<vector<double> > toreturn;
	for (int i = 0; i < count; i++){
		vector<int> tokeep;
		for (int j = 0; j < d.size(); j++) {
			int which =  gsl_rng_uniform(r) * (double) d.size();
			//cout << which << " ";
			tokeep.push_back(which);
		}
		//cout << "\n";
		init_popsizes();
		optim_all_popsizes_bootstrap(tokeep);
		vector<double> tmp;
		for(vector<double>::iterator it = popsizes.begin(); it != popsizes.end(); it++) tmp.push_back(*it);
		toreturn.push_back(tmp);
	}
	return toreturn;
}
void ManyTrees::read_bpoints(string infile){
	bpoints.clear();
	ifstream in(infile.c_str());
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	intStat = stat(infile.c_str(), &stFileInfo);
	if (intStat !=0){
         std::cerr<< "ERROR: cannot open file " << infile << "\n";
         exit(1);
	}
	string st;
	while(getline(in, st)){
		string buf;
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		double bp = atof(line[0].c_str());
		bpoints.push_back(bp);
	}
	vector<double> allcoal = get_all_coaltimes();
	double mc = allcoal[allcoal.size()-1];
	double mb = bpoints[bpoints.size()-1];
	if (mc > mb) bpoints[bpoints.size()-1] = mc+epsilon;
}

void ManyTrees::bootstrap_and_print(string outfile, int count, gsl_rng * r){
	vector<double> means;
	vector<double> ses;

	for(int i = 0; i < popsizes.size(); i++) {
		means.push_back(0);
		ses.push_back(0);
	}

	vector<vector<double> > samps =bootstrap(count, r);
	for(int i = 0; i < samps.size(); i++){
		for (int j = 0; j < samps[i].size(); j++){
			means[j]+= samps[i][j];
		}
	}
	for (int i = 0; i < popsizes.size(); i++) means[i] = means[i]/(double) count;
	for (int i = 0; i < samps.size(); i++){
		for (int j = 0; j < samps[i].size(); j++){
			double m = means[j];
			double diff = samps[i][j] - m;
			ses[j] += diff*diff;
		}
	}
	for (int i = 0; i < popsizes.size(); i++) ses[i] = sqrt(ses[i] /(double) (count-1));
	ofstream out(outfile.c_str());
	for (int i = 0; i < popsizes.size(); i++){
		double t1, t2, p, se;
		t2 = bpoints[i];
		if (i == 0 ) t1 = 0;
		else t1 = bpoints[i-1];
		p = means[i];
		se = ses[i];
		out << t1 << " "<< t2 << " "<< p << " "<< se<< "\n";
	}
}
